import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import torch.nn.functional as F
import torch
from Neural_Networks import (P2D_DSResNet, CNNNet, ProtoNetClassifier, DN4Net, GNNImageClassifier, P2D_DSResNet_GADF,
                             P2D_DSResNet_CWT, P2D_DSResNet_1, P2D_DSResNet_V2)
from train import aLL, Data_loading
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from sklearn.preprocessing import StandardScaler
import numpy as np
from tsne_function import (CNN_extract_features, Proto_extract_features, GNN_extract_features, DN4_extract_features,
                           DSResNet_extract_features, DSResNet_GADF_extract_features, DSResNetCWT_extract_features,
                           P2D_DSResNet_V2_extract_features, P2D_DSResNet_1_extract_features)
import pandas as pd


def cm(modell, test_img, CWT_test_img, true_label):
    n_models = len(modell)
    n_cols = 2  # 每行2列
    n_rows = (n_models + n_cols - 1) // n_cols  # 自动计算行数
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(6 * n_cols, 5 * n_rows))
    axes = axes.flatten()  # 转为一维列表

    if isinstance(axes, plt.Axes):  # 只有一个模型时，axes 不是 list
        axes = [axes]

    for ax, (key, model) in zip(axes, modell.items()):
        model.eval()
        with torch.no_grad():
            if key != 'P2D_DSResNet' and key != 'P2D_DSResNet_V2' and key != 'P2D_DSResNet_1':
                if key == 'P2D_DSResNet_CWT':
                    pre_label = model(CWT_test_img)
                else:
                    pre_label = model(test_img)
            else:
                pre_label = model(test_img, CWT_test_img)
            pre_label = torch.argmax(pre_label, dim=1)

        if isinstance(pre_label, torch.Tensor):
            pre_label = pre_label.cpu().numpy()

        pre_label = pre_label.flatten()
        print(f"{key}:\npre_label:{pre_label.shape}:{pre_label}\ntrue_label:{true_label.shape}:{true_label}\n")
        # 原始混淆矩阵
        cmatrix = confusion_matrix(true_label, pre_label)
        # 按行归一化为百分比
        row_sums = cmatrix.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1  # 避免除以0
        cm_percent = cmatrix.astype('float') / row_sums * 100
        # cm_percent = cmatrix.astype('float') / cmatrix.sum(axis=1, keepdims=True) * 100
        cm_percent = np.round(cm_percent, 2)  # 保留两位小数
        # 显示百分比的文字标签
        disp = ConfusionMatrixDisplay(confusion_matrix=cm_percent)
        disp.plot(ax=ax, cmap="Blues", colorbar=False, values_format=".2f")  # 设置显示格式为百分数

        ax.set_title(f"{key}")
        ax.set_xlabel("Predicted Label")
        ax.set_ylabel("True Label")
    plt.tight_layout()


if __name__ == '__main__':
    # GPU设置
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"dvice:{device}")

    '''数据集分类'''
    # 轴承内圈
    Bearing_inner_ring = [105, 106, 107, 108]
    # 滚动体
    Rolling_element = [118, 119, 120, 121]
    # 外圈
    Outer_ring = [130, 131, 132, 133]
    # 正常数据
    nomal_data = [97]
    data_num = 1
    data1, data_paths = Data_loading()

    num_samples, _, _, _ = aLL(Bearing_inner_ring, Rolling_element, Outer_ring,
                               nomal_data, data_paths, is_cwt=0)
    # 生成固定随机索引（与dataset长度一致）
    np.random.seed(23)
    shared_random_index = np.random.permutation(num_samples)
    # 使用相同随机索引调用 GADF 和 CWT
    _, train_img, train_dl, test_dl = aLL(Bearing_inner_ring, Rolling_element, Outer_ring,
                                          nomal_data, data_paths, is_cwt=1, random_index=shared_random_index)
    _, CWT_train_img, CWT_train_dl, CWT_test_dl = aLL(Bearing_inner_ring, Rolling_element, Outer_ring,
                                                      nomal_data, data_paths, is_cwt=2,
                                                      random_index=shared_random_index)

    all_labels = []
    for _, labels in train_dl:
        all_labels.extend(labels.cpu().numpy())  # 转为 NumPy，放到列表中
    all_labels = np.array(all_labels)
    print('\n')
    print("所有标签数量：", len(all_labels))
    print("标签值：", all_labels)
    print("各类别样本数量：", np.bincount(all_labels))  # 显示每类的数量（假设是从0开始的类别标签）
    print("唯一类别标签：", np.unique(all_labels))
    print('\n')

    '''################### 模型存放 ######################'''
    modell = {}

    # model = CNNNet().to(device)
    # model.load_state_dict(torch.load('0CNNmodel.pth', map_location=device))  # 修改为你保存的文件路径
    # modell['CNN'] = model
    #
    # model2 = ProtoNetClassifier(in_channels=1, hidden_size=64, num_classes=4).to(device)
    # model2.load_state_dict(torch.load('0ProtoNetmodel.pth', map_location=device))  # 修改为你保存的文件路径
    # modell['Proto'] = model2
    #
    model3 = DN4Net(in_channels=1, num_classes=4).to(device)
    model3.load_state_dict(torch.load('0DN4Netmodel.pth', map_location=device))  # 修改为你保存的文件路径
    modell['DN4Net'] = model3
    #
    # model4 = GNNImageClassifier(in_channels=1, num_classes=4).to(device)
    # model4.load_state_dict(torch.load('0GNNmodel.pth', map_location=device))  # 修改为你保存的文件路径
    # modell['GNN'] = model4

    # model5 = P2D_DSResNet(num_classes=4).to(device)
    # model5.load_state_dict(torch.load('0P2D_DSResNetmodel.pth', map_location=device))  # 修改为你保存的文件路径
    # modell['P2D_DSResNet'] = model5

    # model6 = P2D_DSResNet_GADF(num_classes=4).to(device)
    # model6.load_state_dict(torch.load('0P2D_DSResNet_GADFmodel.pth', map_location=device))  # 修改为你保存的文件路径
    # modell['P2D_DSResNet_GADF'] = model6

    # model7 = P2D_DSResNet_CWT(num_classes=4).to(device)
    # model7.load_state_dict(torch.load('0P2D_DSResNet_CWTmodel.pth', map_location=device))  # 修改为你保存的文件路径
    # modell['P2D_DSResNet_CWT'] = model7

    # model8 = P2D_DSResNet_1(num_classes=4).to(device)
    # model8.load_state_dict(torch.load('0P2D_DSResNet_1model.pth', map_location=device))
    # modell['P2D_DSResNet_1'] = model8

    # model9 = P2D_DSResNet_V2(num_classes=4).to(device)
    # model9.load_state_dict(torch.load('0P2D_DSResNet_V2model.pth', map_location=device))
    # modell['P2D_DSResNet_V2'] = model9

    '''###################### tsne图绘制 ######################'''
    # features, labels = CNN_extract_features(model, train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features.numpy())
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels.numpy(), cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of CNN Features')
    # plt.grid(True)
    #
    # features, labels = Proto_extract_features(model2, train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels.numpy(), cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of Proto Features')
    # plt.grid(True)
    #
    features, labels = DN4_extract_features(model3, train_dl, device)
    tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    features_2d = tsne.fit_transform(features)
    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10', alpha=0.7)
    plt.colorbar(scatter, label='Class Label')
    plt.title('t-SNE of DN4 Features')
    plt.grid(True)

    # features, labels = GNN_extract_features(model4, train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of GNN Features')
    # plt.grid(True)

    # features, labels = DSResNet_extract_features(model5, train_dl, CWT_train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # # 创建 DataFrame
    # df = pd.DataFrame({'x': features_2d[:, 0], 'y': features_2d[:, 1]})
    # # 保存为 Excel 文件
    # df.to_excel(r"C:\Users\Lenovo\Desktop\t_sen.xlsx", sheet_name="DSResNet", index=False)  # 不保存行号
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of DSResNet Features')
    # plt.grid(True)

    # features, labels = DSResNet_GADF_extract_features(model6, train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # # 创建 DataFrame
    # df = pd.DataFrame({'x': features_2d[:, 0], 'y': features_2d[:, 1]})
    # with pd.ExcelWriter(r"C:\Users\Lenovo\Desktop\t_sen.xlsx", mode='a', engine='openpyxl') as writer:
    #     df.to_excel(writer, sheet_name='DSResNet_GADF', index=False)
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of DSResNet_GADF Features')
    # plt.grid(True)

    # features, labels = DSResNetCWT_extract_features(model7, train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # # 创建 DataFrame
    # df = pd.DataFrame({'x': features_2d[:, 0], 'y': features_2d[:, 1]})
    # with pd.ExcelWriter(r"C:\Users\Lenovo\Desktop\t_sen.xlsx", mode='a', engine='openpyxl') as writer:
    #     df.to_excel(writer, sheet_name='DSResNetCWT', index=False)
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of DSResNet_CWT Features')
    # plt.grid(True)

    # features, labels = P2D_DSResNet_1_extract_features(model8, train_dl, CWT_train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # # 创建 DataFrame
    # df = pd.DataFrame({'x': features_2d[:, 0], 'y': features_2d[:, 1]})
    # with pd.ExcelWriter(r"C:\Users\Lenovo\Desktop\t_sen.xlsx", mode='a', engine='openpyxl') as writer:
    #     df.to_excel(writer, sheet_name='P2D_DSResNet_1', index=False)
    # plt.figure(figsize=(10, 8))
    # scatter = plt.scatter(features_2d[:, 0], features_2d[:, 1],
    #                       c=labels, cmap='tab10', alpha=0.7)
    # plt.colorbar(scatter, label='Class Label')
    # plt.title('t-SNE of P2D_DSResNet_1 Features')
    # plt.grid(True)

    # features, labels = P2D_DSResNet_V2_extract_features(model9, train_dl, CWT_train_dl, device)
    # tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
    # features_2d = tsne.fit_transform(features)
    # # 创建 DataFrame
    # df = pd.DataFrame({'x': features_2d[:, 0], 'y': features_2d[:, 1]})
    # with pd.ExcelWriter(r"C:\Users\Lenovo\Desktop\t_sen.xlsx", mode='a', engine='openpyxl') as writer:
    #     df.to_excel(writer, sheet_name='P2D_DSResNet', index=False)
    # plt.figure(figsize=(10, 8))
    # sc = plt.scatter(features_2d[:, 0], features_2d[:, 1], c=labels, cmap='tab10', alpha=0.7)
    # plt.colorbar(sc, label='Class Label')
    # plt.title('t-SNE of P2D_DSResNet_V2 Features (after fc1)')
    # plt.grid(True)

    with torch.no_grad():
        i = 0
        for (gadf_img, labels1), (cwt_img, labels2) in zip(train_dl, CWT_train_dl):
            i += 1
            if i == 1:
                assert torch.equal(labels1, labels2), "标签不一致，说明顺序乱了"
                labels = labels1.cpu().numpy()
                print(labels)
                gadf_img = gadf_img.to(device)
                cwt_img = cwt_img.to(device)

                cm(modell, gadf_img, cwt_img, labels)
                break  # 只取一个 batch 作为示例

    # print(i)
    plt.show()
